fix(pt): clone inference coords before enabling grad#5476
Conversation
Avoid enabling gradients in-place on LAMMPS-provided coordinate views. Clone the extended coordinates before building force graphs and pass that cloned tensor to derivative generation while keeping auxiliary metadata private to model output conversion. Closes deepmodeling#5165 Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThree atomic models avoid in-place ChangesCoordinate Gradient Safety and force_coord Propagation
Sequence Diagram(s)sequenceDiagram
participant coord_view as coord_view (leaf view)
participant force_coord as force_coord (clone if needed)
participant atomic as AtomicModel.forward_common_atomic
participant fitter as fit_output_to_model_output
coord_view->>force_coord: if not requires_grad -> detach().clone().requires_grad_(True)
force_coord->>atomic: pass as extended_coord / positional coord
atomic->>fitter: atomic outputs + force_coord -> assemble final model outputs
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Only clone extended coordinates when gradients are not already enabled. This keeps existing autodiff and Hessian graphs intact while still avoiding in-place requires_grad_ on non-grad coordinate views. Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5476 +/- ##
==========================================
- Coverage 82.25% 81.34% -0.92%
==========================================
Files 833 868 +35
Lines 89094 96363 +7269
Branches 4227 4234 +7
==========================================
+ Hits 73287 78387 +5100
- Misses 14515 16676 +2161
- Partials 1292 1300 +8 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
This PR fixes PyTorch inference gradient handling so view tensors passed from lower-level callers such as LAMMPS are not modified in-place when force/virial derivatives are enabled.
Changes:
- Clones coordinate tensors before enabling gradients in PT atomic-model paths.
- Carries the derivative coordinate tensor through
_force_coordmetadata to model output transformation. - Adds regression tests for leaf-view coordinate inputs in atomic and lower model inference paths.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
deepmd/pt/model/atomic_model/base_atomic_model.py |
Preserves private metadata keys while applying atom masks. |
deepmd/pt/model/atomic_model/dp_atomic_model.py |
Clones grad-enabled coordinates and returns _force_coord. |
deepmd/pt/model/atomic_model/linear_atomic_model.py |
Applies the same coordinate cloning and _force_coord propagation for linear models. |
deepmd/pt/model/atomic_model/pairtab_atomic_model.py |
Applies cloning and keeps pairtab energy connected to _force_coord. |
deepmd/pt/model/model/make_model.py |
Uses _force_coord as the autograd derivative source when available. |
source/tests/pt/model/test_dp_atomic_model.py |
Adds atomic-model regression coverage for leaf-view coordinate inputs. |
source/tests/pt/model/test_dp_model.py |
Adds lower-model regression coverage for leaf-view coordinate inputs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| coord = to_torch_tensor(self.coord_ext).requires_grad_(True) | ||
| coord_view = coord.view(self.nf, self.nall, 3) |
| coord_view = coord_ext.requires_grad_(True).view(self.nf, -1, 3) | ||
|
|
||
| ret = md0.forward_lower(coord_view, atype_ext, nlist, do_atomic_virial=True) | ||
|
|
||
| self.assertIn("extended_force", ret) |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt/model/test_dp_atomic_model.py (1)
77-103: ⚡ Quick winAdd documentation and explicit verification for the leaf-view regression test.
This test protects against issue
#5165(RuntimeError when enabling gradients on leaf-view tensors), but that context isn't documented. Future maintainers won't understand why this specific test setup matters.Consider these improvements:
- Add a docstring or comment explaining this is a regression test for issue
#5165and why the leaf-view scenario is critical (LAMMPS provides view tensors)- Add explicit assertions to verify the test preconditions:
self.assertTrue(coord.is_leaf, "coord must be a leaf tensor") self.assertTrue(coord_view._is_view(), "coord_view must be a view")- Consider validating output correctness beyond just key existence—for example, check that energy values are reasonable or match a known reference
The current test will catch crashes (good), but stronger documentation and precondition checks would make the test's purpose clearer and more robust.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/model/test_dp_atomic_model.py` around lines 77 - 103, The test test_forward_common_atomic_accepts_leaf_view_input lacks documentation and explicit precondition checks for the leaf-view regression it protects; add a brief docstring or inline comment referencing issue `#5165` and why a leaf tensor with a view matters, then add explicit assertions before calling md0.forward_* to ensure coord.is_leaf is True and coord_view._is_view() (or coord_view.is_view() depending on API) is True; finally add at least one lightweight correctness check on outputs (e.g., energy tensor shape and finite values from ret["energy"] or a reasonable range) to validate results beyond key existence while keeping the original crash-protection behavior in test_forward_common_atomic_accepts_leaf_view_input.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@source/tests/pt/model/test_dp_atomic_model.py`:
- Around line 77-103: The test
test_forward_common_atomic_accepts_leaf_view_input lacks documentation and
explicit precondition checks for the leaf-view regression it protects; add a
brief docstring or inline comment referencing issue `#5165` and why a leaf tensor
with a view matters, then add explicit assertions before calling md0.forward_*
to ensure coord.is_leaf is True and coord_view._is_view() (or
coord_view.is_view() depending on API) is True; finally add at least one
lightweight correctness check on outputs (e.g., energy tensor shape and finite
values from ret["energy"] or a reasonable range) to validate results beyond key
existence while keeping the original crash-protection behavior in
test_forward_common_atomic_accepts_leaf_view_input.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 35aef69e-ecc3-4b04-8a62-ab185f5ee8c0
📒 Files selected for processing (2)
source/tests/pt/model/test_dp_atomic_model.pysource/tests/pt/model/test_dp_model.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/pt/model/test_dp_model.py
|
Updated the implementation to avoid passing the cloned coordinate through a private This keeps atomic model outputs clean while still avoiding in-place Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5) |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt/model/atomic_model/dp_atomic_model.py`:
- Around line 274-275: In DPAtomicModel.forward_atomic (and any call paths into
forward_common_atomic), avoid calling extended_coord.requires_grad_(True)
in-place on a potential leaf/view tensor; instead create a non-leaf clone of
extended_coord before enabling gradients so autograd/view semantics are
preserved (e.g., clone extended_coord and then set requires_grad on the clone),
and use that cloned tensor for the rest of the forward pass.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d5a0a4a7-fbb5-43e8-af5d-2b26080a3e6c
📒 Files selected for processing (5)
deepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pydeepmd/pt/model/atomic_model/pairtab_atomic_model.pydeepmd/pt/model/model/make_model.pysource/tests/pt/model/test_dp_atomic_model.py
💤 Files with no reviewable changes (1)
- source/tests/pt/model/test_dp_atomic_model.py
Summary
Tests
uvx pre-commit run --files deepmd/pt/model/atomic_model/base_atomic_model.py deepmd/pt/model/atomic_model/dp_atomic_model.py deepmd/pt/model/atomic_model/linear_atomic_model.py deepmd/pt/model/atomic_model/pairtab_atomic_model.py deepmd/pt/model/model/make_model.py source/tests/pt/model/test_dp_atomic_model.py source/tests/pt/model/test_dp_model.pyPYTHONPATH=$PWD uv run --index-strategy unsafe-best-match --with pytest --with numpy --with scipy --with pyyaml --with dargs --with typing_extensions --with h5py --with wcmatch --with packaging --with ml_dtypes --with mendeleev --with array-api-compat --with lmdb --with msgpack --with torch==2.5.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu --no-project pytest source/tests/pt/model/test_dp_model.py::TestDPModel::test_forward_lower_accepts_leaf_view_input source/tests/pt/model/test_dp_atomic_model.py::TestDPAtomicModel::test_forward_common_atomic_accepts_leaf_view_input -qCloses #5165
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
Summary by CodeRabbit
Refactor
Tests